import torch
import torch.nn as nn
import torch.nn.functional as F


class GRN(nn.Module):
    """ 全局响应归一化层（ConvNeXtV2核心组件）"""

    def __init__(self, dim):
        super().__init__()
        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * Nx) + self.beta + x


class ConvNeXtV2Block(nn.Module):
    """ ConvNeXtV2基础块（替换原始残差块）"""

    def __init__(self, dim, layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
        self.norm = nn.LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)
        self.grn = GRN(4 * dim)
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim)) if layer_scale_init_value > 0 else None

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # [N, C, H, W] → [N, H, W, C]
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.grn(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # [N, H, W, C] → [N, C, H, W]
        return input + x


class MultiStreamConvNeXtV2(nn.Module):
    """ 改进后的多流网络（五分类任务）"""

    def __init__(self, in_channels=3, num_classes=5):
        super().__init__()

        # 分支1（原始图像处理）
        self.stem1 = nn.Sequential(
            nn.Conv2d(in_channels, 32, 3, padding=1, bias=False),
            nn.MaxPool2d(2, 2)
        )
        self.block1 = nn.Sequential(
            ConvNeXtV2Block(32),
            ConvNeXtV2Block(32),
            nn.MaxPool2d(2, 2)
        )

        self.stem3 = nn.Sequential(
            nn.Conv2d(32, 64, 3, padding=1, bias=False),
        )

        # 分支2（深度信息处理）
        self.stem2 = nn.Sequential(
            nn.Conv2d(in_channels, 32, 7, padding=3, bias=False),
            nn.MaxPool2d(2, 2)
        )
        self.block2 = nn.Sequential(
            ConvNeXtV2Block(32),
            ConvNeXtV2Block(32),
            nn.MaxPool2d(2, 2)
        )

        # 分支3（结构化数据）
        self.fc_block = nn.Sequential(
            nn.Linear(42, 128),
            nn.BatchNorm1d(128),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512)  # 替换 Dropout 为 BatchNorm1d
        )

        self.class_head = nn.Sequential(
            nn.Linear(512, 1024),
            nn.BatchNorm1d(1024),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(1024, num_classes)
        )

        # 特征融合模块
        self.fusion = nn.ModuleList([
            ConvNeXtV2Block(64),
            ConvNeXtV2Block(128),
            ConvNeXtV2Block(256)
        ])

        # 分类头（ConvNeXtV2设计规范）
        self.before_add = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.LayerNorm(512),
        )

        # 通道调整层
        self.channel_adjust = nn.ModuleDict({
            'layer1': nn.Conv2d(64, 128, 1),
            'layer2': nn.Conv2d(128, 256, 1),
            'layer3': nn.Conv2d(256, 512, 1)
        })

    def forward(self, x1, x2, x3):
        # 分支1处理
        f1 = self.stem1(x1)
        fused = self.block1(f1)
        fused = self.stem3(fused)

        # # 分支2处理
        # f2 = self.stem2(x2)
        # f2 = self.block2(f2)

        # 分支融合
        # fused = torch.cat([f1, f2], dim=1)  # 通道拼接

        # 多级特征融合
        for i, block in enumerate(self.fusion):
            fused = block(fused)
            fused = self.channel_adjust[f'layer{i + 1}'](fused)
            fused = F.max_pool2d(fused, 2)

            # 结构化数据分支
        tabular_feat = self.fc_block(x3)

        # 最终分类
        before_img = self.before_add(fused)
        out = before_img + tabular_feat
        out = self.class_head(out)  # 特征融合

        return out


if __name__ == "__main__":
    model = MultiStreamConvNeXtV2(in_channels=3, num_classes=4)
    rgb = torch.randn(2, 3, 256, 256)
    depth = torch.randn(2, 3, 256, 256)
    tabular = torch.randn(2, 42)
    output = model(rgb, depth, tabular)

    print(output.shape)  # 预期输出: torch.Size([2, 5])